from modelscope import AutoTokenizer
from vllm import LLM, SamplingParams
import os

class LargeLanguageModel:
    def __init__(self, model_name):
        # os.environ["VLLM_LOGGING_LEVEL"] = "ERROR"
        # os.environ["CUDA_VISIBLE_DEVICES"] = "0,1"
        # Initialize the tokenizer
        self.tokenizer = AutoTokenizer.from_pretrained(model_name)

        # Pass the default decoding hyperparameters of Qwen2.5-7B-Instruct
        # max_tokens is for the maximum length for generation.
        self.sampling_params = SamplingParams(temperature=0.1, top_p=0.8, repetition_penalty=1.05, max_tokens=512)

        # Input the model name or path. Can be GPTQ or AWQ models.
        self.llm = LLM(model=model_name, tensor_parallel_size=2)

    def __call__(self, prompt, max_tokens=6):
        # Prepare your prompts
        # prompt = "Tell me something about large language models."
        messages = [
            {"role": "system", "content": "You are Qwen, created by Alibaba Cloud. You are a helpful assistant."},
            {"role": "user", "content": prompt}
        ]
        text = self.tokenizer.apply_chat_template(
            messages,
            tokenize=False,
            add_generation_prompt=True
        )
        self.sampling_params.max_tokens = max_tokens

        # generate outputs
        outputs = self.llm.generate([text], self.sampling_params,use_tqdm=False)

        return outputs[0].outputs[0].text
        # Print the outputs.
        # for output in outputs:
        #     prompt = output.prompt
        #     generated_text = output.outputs[0].text
        #     print(f"Prompt: {prompt!r}, Generated text: {generated_text!r}")